import torch
from torch import nn
import cv2  # pip install opencv-python opencv-contrib-python
import numpy as np

# 映射表
classes = {"0": 0, "1": 1, "2": 2, "3": 3, "4": 4}  # 数字的映射表
# classes = {"0":"Apple","1":"Banana","2":"Watermelon","3":"Pear","4":"Grape"} # 水果映射表
# 加载权重文件
weights = torch.load("./save/best.pt", weights_only=True)
# 将权重文件加载到模型中
model = nn.Sequential(
    nn.Conv2d(3, 16, 3, 1, 1),
    nn.ReLU(),
    nn.MaxPool2d(2, 2),
    nn.Conv2d(16, 32, 3, 1, 1),
    nn.ReLU(),
    nn.MaxPool2d(2, 2),
    nn.Flatten(),
    nn.Dropout(),
    nn.Linear(7 * 7 * 32, 1024),
    nn.ReLU(),
    nn.Linear(1024, 10),
    nn.LogSoftmax(dim=-1)
)
model.load_state_dict(weights)  # 显卡模式下的权重，因此在预测时需要用CUDA
model = model.to(torch.device("cpu"))  # 设置模型为CPU模式（简单推理不需要显卡）

# 加载图片
image = cv2.imread("./data/test/4.png")  # numpy对象  (28x28x3)
image = np.expand_dims(image, 0)  # numpy对象  (1x28x28x3)
image = torch.from_numpy(image)  # torch对象  (1x28x28x3)
image = image.permute([0, 3, 1, 2])  # torch对象  (1x3x28x28)

# 预测结果
model.eval()  # 开启预测模式
predict = model(image.float())
result = torch.argmax(predict, dim=-1)  # 得到分类的下标
print(classes[str(result[0].item())])  # 输出分类的名称
